package multi_locks

import (
	"sync"
)

type MultiRWMux struct {
	n   int64
	mux []sync.RWMutex
}

func NewMultiRWMux(n int64) *MultiRWMux {
	if n <= 0 {
		n = 1
	}

	mm := &MultiRWMux{
		n:   n,
		mux: make([]sync.RWMutex, n),
	}

	return mm
}

func (mm *MultiRWMux) GetMuxByUint32(u32 uint32) *sync.RWMutex {
	if mm.n == 1 {
		return &mm.mux[0]
	}

	idx := int64(u32) % mm.n
	return &mm.mux[idx]
}

func (mm *MultiRWMux) GetMuxByUint64(u64 uint64) *sync.RWMutex {
	if mm.n == 1 {
		return &mm.mux[0]
	}

	idx := u64 % uint64(mm.n)
	return &mm.mux[idx]
}
